/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package jpsearch;

import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;

import org.apache.sanselan.ImageReadException;
import org.apache.sanselan.ImageWriteException;
import org.apache.sanselan.common.BinaryFileParser;
import org.apache.sanselan.common.byteSources.ByteSource;
import org.apache.sanselan.formats.jpeg.JpegUtils;

/**
 * Interface for Exif write/update/remove functionality for Jpeg/JFIF images.
 * <p>
 * <p>
 * 
 */
public class JpegRewriter extends BinaryFileParser implements JPSearchConstants
{
	private static final int JPEG_BYTE_ORDER = BYTE_ORDER_NETWORK;

	/**
	 * Constructor. to guess whether a file contains an image based on its file
	 * extension.
	 */
	public JpegRewriter()
	{
		setByteOrder(JPEG_BYTE_ORDER);
	}

	protected static class JFIFPieces
	{
		public final List pieces;
		public final List segmentPieces;

		public JFIFPieces(final List pieces, final List segmentPieces)
		{
			this.pieces = pieces;
			this.segmentPieces = segmentPieces;
		}

	}

	protected abstract static class JFIFPiece
	{
		protected abstract void write(OutputStream os) throws IOException;
		
		public String toString()
		{
			return "[" + this.getClass().getName() + "]";
		}
	}

	protected static class JFIFPieceSegment extends JFIFPiece
	{
		public final int marker;
		public final byte markerBytes[];
		public final byte segmentLengthBytes[];
		public final byte segmentData[];

		public JFIFPieceSegment(final int marker, final byte[] segmentData)
		{
			this(marker, int2ToByteArray(marker, JPEG_BYTE_ORDER),
					int2ToByteArray(segmentData.length + 2, JPEG_BYTE_ORDER),
					segmentData);
		}

		public JFIFPieceSegment(final int marker, final byte[] markerBytes,
				final byte[] segmentLengthBytes, final byte[] segmentData)
		{
			this.marker = marker;
			this.markerBytes = markerBytes;
			this.segmentLengthBytes = segmentLengthBytes;
			this.segmentData = segmentData;
		}

		public String toString()
		{
			return "[" + this.getClass().getName() + " (0x" + Integer.toHexString(marker) + ")]";
		}

		protected void write(OutputStream os) throws IOException
		{
			os.write(markerBytes);
			os.write(segmentLengthBytes);
			os.write(segmentData);
		}

		public boolean isApp1Segment()
		{
			return marker == JPEG_APP1_Marker;
		}

		public boolean isAppSegment()
		{
			return marker >= JPEG_APP0_Marker && marker <= JPEG_APP15_Marker;
		}

		public boolean isExifSegment()
		{
			if (marker != JPEG_APP1_Marker)
				return false;
			if (!byteArrayHasPrefix(segmentData, EXIF_IDENTIFIER_CODE))
				return false;
			return true;
		}

		public boolean isApp3Segment()
		{
			return marker == JPEG_APP3_Marker;
		}

	}

	protected static class JFIFPieceImageData extends JFIFPiece
	{
		public final byte markerBytes[];
		public final byte imageData[];

		public JFIFPieceImageData(final byte[] markerBytes,
				final byte[] imageData)
		{
			super();
			this.markerBytes = markerBytes;
			this.imageData = imageData;
		}

		protected void write(OutputStream os) throws IOException
		{
			os.write(markerBytes);
			os.write(imageData);
		}
	}

	protected JFIFPieces analyzeJFIF(ByteSource byteSource)
			throws ImageReadException, IOException
	// , ImageWriteException
	{
		final ArrayList pieces = new ArrayList();
		final List segmentPieces = new ArrayList();

		JpegUtils.Visitor visitor = new JpegUtils.Visitor() {
			// return false to exit before reading image data.
			public boolean beginSOS()
			{
				return true;
			}

			public void visitSOS(int marker, byte markerBytes[],
					byte imageData[])
			{
				pieces.add(new JFIFPieceImageData(markerBytes, imageData));
			}

			// return false to exit traversal.
			public boolean visitSegment(int marker, byte markerBytes[],
					int segmentLength, byte segmentLengthBytes[],
					byte segmentData[]) throws ImageReadException, IOException
			{
				JFIFPiece piece = new JFIFPieceSegment(marker, markerBytes,
						segmentLengthBytes, segmentData);
				pieces.add(piece);
				segmentPieces.add(piece);

				return true;
			}
		};

		new JpegUtils().traverseJFIF(byteSource, visitor);

		return new JFIFPieces(pieces, segmentPieces);
	}

	private static interface SegmentFilter
	{
		public boolean filter(JFIFPieceSegment segment);
	}

	private static final SegmentFilter EXIF_SEGMENT_FILTER = new SegmentFilter() {
		public boolean filter(JFIFPieceSegment segment)
		{
			return segment.isExifSegment();
		}
	};

	private static final SegmentFilter JPSEARCH_SEGMENT_FILTER = new SegmentFilter() {
		public boolean filter(JFIFPieceSegment segment)
		{
			return segment.isApp3Segment();
		}
	};
	

	protected List removeJPSearchSegments(List segments)
	{
		return filterSegments(segments, JPSEARCH_SEGMENT_FILTER);
	}

	protected List removeExifSegments(List segments)
	{
		return filterSegments(segments, EXIF_SEGMENT_FILTER);
	}

	protected List filterSegments(List segments, SegmentFilter filter)
	{
		return filterSegments(segments, filter, false);
	}

	protected List filterSegments(List segments, SegmentFilter filter,
			boolean reverse)
	{
		List result = new ArrayList();

		for (int i = 0; i < segments.size(); i++)
		{
			JFIFPiece piece = (JFIFPiece) segments.get(i);
			if (piece instanceof JFIFPieceSegment)
			{
				if (filter.filter((JFIFPieceSegment) piece) ^ !reverse)
					result.add(piece);
			} else if(!reverse)
				result.add(piece);
		}

		return result;
	}

	protected List insertBeforeFirstAppSegments(List segments, List newSegments)
			throws ImageWriteException
	{
		int firstAppIndex = -1;
		for (int i = 0; i < segments.size(); i++)
		{
			JFIFPiece piece = (JFIFPiece) segments.get(i);
			if (!(piece instanceof JFIFPieceSegment))
				continue;

			JFIFPieceSegment segment = (JFIFPieceSegment) piece;
			if (segment.isAppSegment())
			{
				if (firstAppIndex == -1)
					firstAppIndex = i;
			}
		}

		List result = new ArrayList(segments);
		if (firstAppIndex == -1)
			throw new ImageWriteException("JPEG file has no APP segments.");
		result.addAll(firstAppIndex, newSegments);
		return result;
	}

	protected List insertAfterLastAppSegments(List segments, List newSegments)
			throws ImageWriteException
	{
		int lastAppIndex = -1;
		for (int i = 0; i < segments.size(); i++)
		{
			JFIFPiece piece = (JFIFPiece) segments.get(i);
			if (!(piece instanceof JFIFPieceSegment))
				continue;

			JFIFPieceSegment segment = (JFIFPieceSegment) piece;
			if (segment.isAppSegment())
				lastAppIndex = i;
		}

		List result = new ArrayList(segments);
		if (lastAppIndex == -1)
		{
			if(segments.size()<1)
				throw new ImageWriteException("JPEG file has no APP segments.");
			result.addAll(1, newSegments);
		}
		else
		result.addAll(lastAppIndex + 1, newSegments);
		
		return result;
	}

	protected void writeSegments(OutputStream os, List segments)
			throws ImageWriteException, IOException
	{
		try
		{
			os.write(SOI);

			for (int i = 0; i < segments.size(); i++)
			{
				JFIFPiece piece = (JFIFPiece) segments.get(i);
				piece.write(os);
			}
			os.close();
			os = null;
		} finally
		{
			try
			{
				if (os != null)
					os.close();
			} catch (Exception e)
			{
				// swallow exception; already in the context of an exception.
			}
		}
	}


	public static class JpegSegmentOverflowException extends
			ImageWriteException
	{
		public JpegSegmentOverflowException(String s)
		{
			super(s);
		}
	}

}